import os

import numpy as np
from PIL import Image
from dataset import *
from glob import glob
from tqdm import tqdm

cad_load_path = os.path.join('.', 'data', 'shrec14-views')
sk_load_path = os.path.join('.', 'data', 'shrec14-sketch')
# cad_load_path = os.path.join('.', 'data', 'shrec13-views')
# sk_load_path = os.path.join('.', 'data', 'shrec13-sketch')


def getMeanAndStd(dataset: str):
    """Calculate the mean and standard deviation of the data set
    
    Args:
        dataset: 'sketch' or 'views'.

    """
    img_list = list()
    image_fns = list()
    if dataset == 'all':
        image_fns = glob(os.path.join(cad_load_path, '*', '*', '*')) + glob(os.path.join(sk_load_path, '*', '*', '*'))
    elif dataset == 'views':
        image_fns = glob(os.path.join(cad_load_path, '*', '*', '*'))
    else:
        image_fns = glob(os.path.join(sk_load_path, '*', '*', '*'))
    for path in tqdm(image_fns):
        img = np.array(Image.open(path))
        img = img[:, :, np.newaxis]
        img_list.append(img)

    imgs = np.concatenate(img_list, axis=2)
    imgs = imgs.astype(np.float32) / 255.
    print('Shape: ({}, {}), total: {}'.format(imgs.shape[0], imgs.shape[1], imgs.shape[2]))

    pixels = imgs.ravel()
    means = np.mean(pixels)
    stdevs = np.std(pixels)
    print('Normalize Mean of {} = {}'.format(dataset, means))
    print('Normalize Std of {} = {}'.format(dataset, stdevs))


def getFeaturesMeanAndStd():
    feature_train_data = FeatureDataset(feature_path, 'train')
    feature_test_data = FeatureDataset(feature_path, 'test')
    feature_train_loader = DataLoader(dataset=feature_train_data, batch_size=64, shuffle=True)
    feature_test_loader = DataLoader(dataset=feature_test_data, batch_size=64, shuffle=True)

    feature_list = []
    for _, data in enumerate(tqdm(feature_train_loader)):
        f, _, _ = data
        feature_list.append(f)
    for _, data in enumerate(tqdm(feature_test_loader)):
        f, _, _ = data
        feature_list.append(f)

    feature_list = np.concatenate(feature_list)
    feature_list = feature_list.ravel()
    means = np.mean(feature_list)
    stdevs = np.std(feature_list)
    maxv = np.max(feature_list)
    minv = np.min(feature_list)
    print('Normalize Mean of features = {}'.format(means))
    print('Normalize Std of features = {}'.format(stdevs))
    print('Max and min value of features = {}, {}'.format(maxv, minv))


if __name__ == '__main__':
    # part-shrec14
    #   All:  mean = 0.9863767623901367, std = 0.11683730781078339
    #   Views: mean = 0.9675745964050293, std = 0.18061445653438568
    #   Sketch: mean = 0.9791669249534607, std = 0.10922393947839737
    # shrec14:
    #   All: mean = 0.9693331718444824, std = 0.1738644391298294
    #   Views: mean = 0.9675745964050293, std = 0.18061445653438568
    #   Shapes: mean = 0.9731981754302979, std = 0.11050379276275635
    #   Sketch: mean = 0.9804596304893494, std = 0.10598362982273102
    # shrec13:
    #   All: mean = 0.9839804768562317, std = 0.11487835645675659
    #   Views: mean = 0.9859437942504883, std = 0.1182224228978157
    #   Sketch: mean = 0.9798703193664551, std = 0.10752718895673752
    getMeanAndStd('all')
    getMeanAndStd('views')
    getMeanAndStd('sketch')
    # part-shrec14 features
    #   Mean of features = 0.7377912402153015, Std of features = 0.2630916237831116
    #   Max of features = 3.2473697662353516
    #   Min of features = -0.07844455540180206
    # getFeaturesMeanAndStd()
